#include <iostream>
#include <eigen3/Eigen/Dense>
#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <nav_msgs/Odometry.h>
#include <sensor_msgs/PointCloud2.h>
#include <pcl_conversions/pcl_conversions.h>
#include <string>
#include <vector>
#include <XmlRpcValue.h>
#include <unordered_map>
#include <traj_utils/DetectPose.h>

using std::cout;
using std::endl;
using std::vector;

/* -------------------------------------------------------------------------- */
/*                              class definition                              */
/* -------------------------------------------------------------------------- */
class VoxelGrid
{
public:
  Eigen::Vector3d map_origin_, map_size_;
  Eigen::Vector3i map_voxel_num_;
  double resolution_, resolution_inv_;
  double obstacles_inflation_;
  double ground_height_;

  std::vector<char> occupancy_buffer_inflate_;

  void init(ros::NodeHandle &nh)
  {
    /* get parameter */
    double x_size, y_size, z_size;
    nh.param("VoxelGrid/resolution", resolution_, -1.0);
    nh.param("VoxelGrid/map_size_x", x_size, -1.0);
    nh.param("VoxelGrid/map_size_y", y_size, -1.0);
    nh.param("VoxelGrid/map_size_z", z_size, -1.0);
    nh.param("VoxelGrid/obstacles_inflation", obstacles_inflation_, -1.0);
    nh.param("VoxelGrid/ground_height", ground_height_, 1.0);

    resolution_inv_ = 1 / resolution_;
    map_origin_ = Eigen::Vector3d(-x_size / 2.0, -y_size / 2.0, ground_height_);
    map_size_ = Eigen::Vector3d(x_size, y_size, z_size);
    for (int i = 0; i < 3; ++i)
      map_voxel_num_(i) = ceil(map_size_(i) / resolution_);

    int buffer_size = map_voxel_num_(0) * map_voxel_num_(1) * map_voxel_num_(2);
    occupancy_buffer_inflate_ = std::vector<char>(buffer_size, 0);
  }

  void posToIndex(const Eigen::Vector3d &pos, Eigen::Vector3i &id)
  {
    for (int i = 0; i < 3; ++i)
      id(i) = floor((pos(i) - map_origin_(i)) * resolution_inv_);
  }

  int toAddress(const Eigen::Vector3i &id)
  {
    return id(0) * map_voxel_num_(1) * map_voxel_num_(2) +
           id(1) * map_voxel_num_(2) + id(2);
  }

  bool isInMap(const Eigen::Vector3i &idx)
  {
    if (idx(0) < 0 || idx(1) < 0 || idx(2) < 0)
    {
      return false;
    }
    if (idx(0) > map_voxel_num_(0) - 1 ||
        idx(1) > map_voxel_num_(1) - 1 ||
        idx(2) > map_voxel_num_(2) - 1)
    {
      return false;
    }
    return true;
  }

  void setInput(pcl::PointCloud<pcl::PointXYZ> cloud)
  {
    int inf_step = ceil(obstacles_inflation_ / resolution_);
    pcl::PointXYZ pt;
    Eigen::Vector3d p3d, p3d_inf;

    for (size_t i = 0; i < cloud.points.size(); ++i)
    {
      pt = cloud.points[i];
      p3d(0) = pt.x, p3d(1) = pt.y, p3d(2) = pt.z;

      /* point inside update range */
      Eigen::Vector3i index;

      /* inflate the point */
      for (int x = -inf_step; x <= inf_step; ++x)
        for (int y = -inf_step; y <= inf_step; ++y)
          for (int z = -inf_step; z <= inf_step; ++z)
          {

            p3d_inf(0) = pt.x + x * resolution_;
            p3d_inf(1) = pt.y + y * resolution_;
            p3d_inf(2) = pt.z + z * resolution_;

            posToIndex(p3d_inf, index);

            if (!isInMap(index))
              continue;

            int idx_inf = toAddress(index);

            occupancy_buffer_inflate_[idx_inf] = 1;
          }
    }
  }

  bool isBlocked(const Eigen::Vector3d &start, const Eigen::Vector3d &end)
  {
    Eigen::Vector3d dir = end - start;
    double length = dir.norm();
    dir.normalize();
    Eigen::Vector3d pos = start;
    for (double i = 0; i < length; i += resolution_)
    {
      pos += dir * resolution_;
      Eigen::Vector3i idx;
      posToIndex(pos, idx);
      if (!isInMap(idx))
        return true;
      int addr = toAddress(idx);
      if (occupancy_buffer_inflate_[addr] == 1)
        return true;
    }
    return false;
  }
};

// 同事订阅多个话题，并将消息存储到一个map中
class DroneOdomListener
{
public:
  DroneOdomListener() {}
  DroneOdomListener(ros::NodeHandle &nh,
                    const std::vector<std::string> &topic_names)
  {
    for (const std::string &topic : topic_names) {
      // 为每个话题创建一个订阅者，并将回调函数绑定到相应的话题上
      ros::Subscriber sub =
          nh.subscribe<nav_msgs::Odometry>(
              topic, 1,
              boost::bind(&DroneOdomListener::odomCallback,
                          this, _1, topic));
      subscribers_.push_back(sub);

      int id = extractIdFromTopicName(topic);
      ids_.push_back(id);
    }
  }
  ~DroneOdomListener() {}

  std::unordered_map<int, geometry_msgs::PoseStamped> getPosMap()
  {
    return pos_map_;
  }

  void odomCallback(const nav_msgs::Odometry::ConstPtr &msg,
                    const std::string &topic_name)
  {
    // 从话题名称中解析出ID
    int id = extractIdFromTopicName(topic_name);
    // 根据ID存储消息
    geometry_msgs::PoseStamped pose;
    pose.header = msg->header;
    pose.pose = msg->pose.pose;

    pos_map_[id] = pose;
  }

  typedef std::shared_ptr<DroneOdomListener> Ptr;

private:
  std::vector<ros::Subscriber> subscribers_;
  std::unordered_map<int, geometry_msgs::PoseStamped> pos_map_;
  std::vector<int> ids_;

  int extractIdFromTopicName(const std::string &topic_name)
  {
    // 从话题名中解析ID，这里假设格式总是 "/drone_X_visual_slam/odom"
    std::size_t id_start = topic_name.find("_") + 1;
    // std::size_t id_end = topic_name.find("_", id_start);
    // std::string id_str = topic_name.substr(id_start, id_end - id_start);
    std::string id_str = topic_name.substr(id_start, 1);
    return std::stoi(id_str);
  }
};

/* -------------------------------------------------------------------------- */
/*                              Global Variables                              */
/* -------------------------------------------------------------------------- */
ros::Subscriber cam_pose_sub, global_map_sub;
ros::Publisher detect_pub;
ros::Timer detect_timer;

double detect_freq = 10.0;

geometry_msgs::PoseStamped camera_pose;

VoxelGrid global_map;
DroneOdomListener::Ptr listener;

bool flag_have_camera = false;
bool flag_have_global_map = false;

/* configuration */
// FOV = 2 * atan2(0.5 * image_width, focal_length)
double max_sight_length = 7.0;  // TODO: set this value
double cosVerticalFOV = 0.849;  // TODO: set this value
double cosHorizontalFOV = 0.77; // TODO: set this value

std::unordered_map<int, Eigen::Vector4d> pos_record;

/* -------------------------------------------------------------------------- */
/*                                  tool func                                 */
/* -------------------------------------------------------------------------- */

void getStringVectorParam(ros::NodeHandle &nh,
                          const std::string &param_name,
                          std::vector<std::string> &str_vec)
{
  XmlRpc::XmlRpcValue xml_array;
  if (nh.getParam(param_name, xml_array) &&
      xml_array.getType() == XmlRpc::XmlRpcValue::TypeArray)
  {
    for (int i = 0; i < xml_array.size(); ++i)
    {
      if (xml_array[i].getType() == XmlRpc::XmlRpcValue::TypeString)
      {
        str_vec.push_back(static_cast<std::string>(xml_array[i]));
      }
      else
      {
        ROS_ERROR("Non-string element found in string_array_param; ignoring.");
      }
    }
  }
  else
  {
    ROS_ERROR("Failed to get param '%s' or it is not an array.",
              param_name.c_str());
  }

  for (const std::string &str : str_vec)
  {
    ROS_INFO("Array element: %s", str.c_str());
  }
}

/* -------------------------------------------------------------------------- */
/*                              callback function                             */
/* -------------------------------------------------------------------------- */
void detectCb(const ros::TimerEvent &event)
{
  if (!flag_have_global_map)
    return;
  if (!flag_have_camera)
    return;

  geometry_msgs::Point zeros;
  zeros.x = 0;
  zeros.y = 0;
  zeros.z = 0;

  std::unordered_map<int, geometry_msgs::PoseStamped>
      pos_map = listener->getPosMap();

  traj_utils::DetectPose detect_result;

  for (std::unordered_map<int, geometry_msgs::PoseStamped>::iterator it = pos_map.begin();
       it != pos_map.end(); ++it)
  {
    geometry_msgs::PoseStamped target_pose = it->second;

    if (abs((camera_pose.header.stamp - target_pose.header.stamp).toSec()) > 0.1)
      continue;

    Eigen::Vector3d self(camera_pose.pose.position.x,
                         camera_pose.pose.position.y,
                         camera_pose.pose.position.z);
    Eigen::Vector3d target(target_pose.pose.position.x,
                           target_pose.pose.position.y,
                           target_pose.pose.position.z);

    Eigen::Vector3d V = target - self;
    if (V.norm() > max_sight_length)
      continue;

    Eigen::Quaterniond quad(camera_pose.pose.orientation.w,
                            camera_pose.pose.orientation.x,
                            camera_pose.pose.orientation.y,
                            camera_pose.pose.orientation.z);
    V = quad.inverse() * V;

    // check if the target is in the camera's field of view
    // XZ平面投影，忽略Y分量
    Eigen::Vector3d Vh = V;
    Vh.y() = 0;
    double cosAlphaH = Vh.normalized().dot(Eigen::Vector3d(0, 0, 1));
    if (cosAlphaH < cosHorizontalFOV)
      continue;

    // YZ平面投影，忽略X分量
    Eigen::Vector3d Vv = V;
    Vv.x() = 0;
    double cosAlphaV = Vv.normalized().dot(Eigen::Vector3d(0, 0, 1));
    if (cosAlphaV < cosVerticalFOV)
      continue;

    // check if the line is blocked
    if (global_map.isBlocked(target, self))
      continue;

    detect_result.ids.push_back(it->first);
    detect_result.poses.push_back(target_pose.pose.position);

    geometry_msgs::Point vel;
    if (pos_record.find(it->first) != pos_record.end())
    {
      double delta_t = (target_pose.header.stamp).toSec() - pos_record[it->first](3);
      vel.x = (target[0] - pos_record[it->first](0)) / delta_t;
      vel.y = (target[1] - pos_record[it->first](1)) / delta_t;
      vel.z = (target[2] - pos_record[it->first](2)) / delta_t;
    }
    else {
      vel = zeros;
    }
    detect_result.vels.push_back(vel);

    Eigen::Vector4d ptime(target[0], target[1], target[2],
                          target_pose.header.stamp.toSec());
    pos_record[it->first] = ptime;
  }

  if (detect_result.ids.size() == 0) return;

  detect_result.header = camera_pose.header;
  detect_pub.publish(detect_result);
}

void CamPoseCb(const geometry_msgs::PoseStampedConstPtr &msg)
{
  camera_pose = *msg;
  flag_have_camera = true;
}

void rcvGlobalPointCloudCb(const sensor_msgs::PointCloud2 &pointcloud_map)
{
  if (flag_have_global_map)
    return;

  // load global map
  pcl::PointCloud<pcl::PointXYZ> cloudIn;

  // transform map to point cloud format
  pcl::fromROSMsg(pointcloud_map, cloudIn);

  global_map.setInput(cloudIn);
  ROS_WARN("Fake drone detect: global map received.");

  flag_have_global_map = true;
}

/* -------------------------------------------------------------------------- */
/*                                  main func                                 */
/* -------------------------------------------------------------------------- */

int main(int argc, char **argv)
{
  ros::init(argc, argv, "fake_drone_detect");
  ros::NodeHandle nh("~");

  nh.param("DroneDetect/frequent", detect_freq, 10.0);
  nh.param("DroneDetect/max_sight_length", max_sight_length, 7.0);

  global_map.init(nh);

  std::vector<std::string> target_odom_array;
  getStringVectorParam(nh, "/targer_pose_array", target_odom_array);

  listener = std::make_shared<DroneOdomListener>(nh, target_odom_array);

  detect_pub = nh.advertise<traj_utils::DetectPose>("detect_result", 1);

  cam_pose_sub = nh.subscribe("camera_pose", 1, CamPoseCb);
  global_map_sub = nh.subscribe("global_map", 1, rcvGlobalPointCloudCb);

  detect_timer = nh.createTimer(ros::Duration(1 / detect_freq), detectCb);

  ros::spin();

  return 0;
}